import pandas as pd
from peewee import *

from service.das.database_pool import DatabasePool


# 周K、月K、年K数据表定义
class KLineData(Model):
    id = AutoField(primary_key=True)
    stock_code = CharField(index=True)
    trade_date = DateTimeField(index=True)  # 聚合到的时间点
    open = FloatField()
    high = FloatField()
    low = FloatField()
    close = FloatField()
    pre_close = FloatField(null=True)
    change = FloatField(null=True)
    pct_chg = FloatField(null=True)
    vol = FloatField(null=True)
    amount = FloatField(null=True)
    type = CharField(choices=[("W", "周K"), ("M", "月K"), ("Y", "年K")])  # 区分类型

    class Meta:
        database = DatabasePool(max_connections=5).database
        table_name = "kline_data"


class StockKLineDas:
    # 新增K线数据
    def add_kline_data(data):
        """
        新增一条K线数据
        :param data: 字典类型，包含K线数据字段
        """
        KLineData.create(**data)

    # 查询K线数据
    def get_kline_data(stock_code, kline_type, start_date=None, end_date=None):
        """
        查询指定条件的K线数据
        :param stock_code: 股票代码
        :param kline_type: K线类型（'W', 'M', 'Y'）
        :param start_date: 起始日期
        :param end_date: 结束日期
        """
        query = KLineData.select().where(
            KLineData.stock_code == stock_code, KLineData.type == kline_type
        )
        if start_date:
            query = query.where(KLineData.trade_date >= start_date)
        if end_date:
            query = query.where(KLineData.trade_date <= end_date)
        return list(query.dicts())

    # 更新K线数据
    def update_kline_data(kline_id, update_fields):
        """
        更新一条K线数据
        :param kline_id: K线数据的ID
        :param update_fields: 更新的字段和值
        """
        KLineData.update(**update_fields).where(KLineData.id == kline_id).execute()

    # 删除K线数据
    def delete_kline_data(kline_id):
        """
        删除一条K线数据
        :param kline_id: K线数据的ID
        """
        KLineData.delete().where(KLineData.id == kline_id).execute()

    def delete_kline_data(
        self, stock_code: Optional[str] = None, kline_type: Optional[str] = None
    ):
        """
        删除指定股票的所有 K 线数据，支持按类型筛选。

        Args:
            stock_code (Optional[str]): 股票代码。如果为 None，则删除所有股票的数据。
            kline_type (Optional[str]): K 线类型 ('W', 'M', 'Y')。如果为 None，则删除所有类型的数据。

        Returns:
            int: 删除的记录数。
        """
        query = KLineData.delete()

        # 根据 stock_code 筛选
        if stock_code:
            query = query.where(KLineData.stock_code == stock_code)

        # 根据 K 线类型筛选
        if kline_type:
            query = query.where(KLineData.type == kline_type)

        # 执行删除操作
        deleted_count = query.execute()
        AntLogger.info
            f"Deleted {deleted_count} records for stock_code={stock_code}, kline_type={kline_type}."
        )
        return deleted_count

    def store_aggregated_kline(self, aggregated_data: pd.DataFrame, kline_type: str):
        """
        将聚合后的 K 线数据存储到目标表。
        """
        for _, row in aggregated_data.iterrows():
            KLineData.create(
                stock_code=row["stock_code"],
                trade_date=row["group"],
                open=row["open"],
                high=row["high"],
                low=row["low"],
                close=row["close"],
                vol=row["vol"],
                amount=row["amount"],
                type=kline_type,
            )
